
import gc
import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
from pickle import dump

from constant import OUTPUT_PATH,TESTB_PATH,TRAIN_PATH,\
    DatasetLabelPath,DatasetTrainPath,DatasetTestPath,\
        LabelScalerPath,DataScalerPath

train=pd.read_csv(TRAIN_PATH, sep=' ')
testB=pd.read_csv(TESTB_PATH, sep=' ')
train['train']=1
testB['train']=0
dataset=pd.concat([train,testB],ignore_index=True)
del train,testB
gc.collect()

# 处理特征的非数值取值
dataset.replace({'notRepairedDamage':'-'}, 2, inplace=True)
dataset['notRepairedDamage']=dataset['notRepairedDamage'].map(float).astype(np.int8)

# 时间数据类型转换
dataset['regDate']=pd.to_datetime(dataset['regDate'],format="%Y%m%d",errors='coerce')
dataset['creatDate']=pd.to_datetime(dataset['creatDate'],format="%Y%m%d",errors='coerce')
dataset.drop(labels=['offerType','seller'],axis=1,inplace=True)

dataset.drop(['SaleID','name'],axis=1,inplace=True)
# 题目要求power的范围是[0,600]
# 此处进行截断
dataset['power']=dataset['power'].map(lambda x: x if x<=600 else 600)

# 处理空值
dataset.fillna({'regDate':dataset['regDate'].median()},inplace=True)
dataset.dropna(subset=['model'],inplace=True)
dataset.fillna({'bodyType':dataset['bodyType'].median()},inplace=True)
dataset.fillna({'fuelType':dataset['fuelType'].median()},inplace=True)
dataset.fillna({'gearbox':dataset['gearbox'].median()},inplace=True)

# 获取时间特征
dataset['usedtime']=(dataset['creatDate']-dataset['regDate']).map(lambda x:x.days)
dataset.drop(['creatDate','regDate'],axis=1,inplace=True)

# 调整数据类型
for key in ['model', 'bodyType', 'fuelType','gearbox']:
    dataset[key] = dataset[key].astype(np.int64)

# 保存数据集
dataset.to_csv(OUTPUT_PATH+'dataset.csv', index=False)

# 划分数据
datasetLabel = dataset[dataset['train']==1].drop(['train'],axis=1)[['price']].values
label_scaler = MinMaxScaler().fit(datasetLabel)
with open(LabelScalerPath,'bw') as f:
    dump(label_scaler, f)
datasetLabel = label_scaler.transform(datasetLabel)
np.save(DatasetLabelPath, datasetLabel)
print(datasetLabel.shape)

datasetTrain = dataset[dataset['train']==1].drop(['train','price'],axis=1).values
data_scaler = MinMaxScaler().fit(datasetTrain)
with open(DataScalerPath,'bw') as f:
    dump(data_scaler, f)
datasetTrain = data_scaler.transform(datasetTrain)
np.save(DatasetTrainPath, datasetTrain)
print(datasetTrain.shape)

datasetTest=dataset[dataset['train']==0].drop(['train','price'],axis=1).values
datasetTest = data_scaler.transform(datasetTest)
np.save(DatasetTestPath, datasetTest)
print(datasetTest.shape)

del dataset
gc.collect()